#pragma once
#include <TlHelp32.h>

struct PAddrType
{
	const wxString name;
	const wxString fmt;
	const u32 size;
};

static const PAddrType m_types[] =
{
	{"8 bit", "%d", sizeof(u8)},
	{"16 bit", "%d", sizeof(u16)},
	{"32 bit", "%d", sizeof(u32)},
	{"64 bit", "%lld", sizeof(u64)},
	
	{"float", "%f", sizeof(float)},
	{"double", "%f", sizeof(double)},
};

static const u32 m_types_count = WXSIZEOF(m_types);

struct PAddr
{
	const PAddrType& type;
	const u64 type_id;
	u8* addr;
	u64 value;
	u32 size;

	PAddr(const PAddrType& type, const u64 type_id, u8* addr, u32 size)
		: type(type)
		, type_id(type_id)
		, addr(addr)
		, value(0)
		, size(size)
	{
	}

	~PAddr()
	{
	}
};

class ProcessManager
{
	Array<PROCESSENTRY32> m_process_list;
	u32 m_process;
	HANDLE m_proc_handle;

	struct MemoryPage
	{
		u8* addr;
		u32 size;
	};

	Array<MemoryPage> m_proc_pages;
	u32 m_page_size;

public:
	ProcessManager()
	{
		Update();
	}

	static u64 re(const u64 data, u8 size)
	{
		/*
		switch(size)
		{
		case sizeof(u16): return _byteswap_ushort(data);
		case sizeof(u32): return _byteswap_ulong(data);
		case sizeof(u64): return _byteswap_uint64(data);
		}
		*/

		return data;
	}

	static bool GetDebugPrivilege()
	{
		enum { SeDebugPrivilege = 20 };

		LONG (WINAPI *RtlAdjustPrivilege)(DWORD, BOOL, INT, PBOOL);
		*(FARPROC*)&RtlAdjustPrivilege = GetProcAddress(GetModuleHandle("ntdll.dll"), "RtlAdjustPrivilege");

		if(!RtlAdjustPrivilege) return false;

		BOOL last_value = FALSE;
		RtlAdjustPrivilege(SeDebugPrivilege, TRUE, 0, &last_value);

		return true;
	}

	void Update()
	{
		m_process_list.Clear();

		PROCESSENTRY32 pe = {sizeof(PROCESSENTRY32)};
		u64 pid = 0;

		if(HANDLE ps = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0))
		{
			for(BOOL cond = Process32First(ps, &pe); cond; cond = Process32Next(ps, &pe))
			{
				m_process_list.AddCpy(pe);
			}

			CloseHandle(ps);
		}
	}

	bool Open(const wxString& process)
	{
		m_proc_pages.Clear();

		const PROCESSENTRY32* proc;
		if(!(proc = GetProcess(process))) return false;

		m_proc_handle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, proc->th32ProcessID);

		if(!m_proc_handle) return false;

		MEMORY_BASIC_INFORMATION mbi;
		memset(&mbi, 0, sizeof(mbi));

		SYSTEM_INFO msi;
		GetSystemInfo(&msi);

		for(u8* addr = (u8*)msi.lpMinimumApplicationAddress; addr <= (u8*)msi.lpMaximumApplicationAddress; addr += mbi.RegionSize)
		{
			if(!VirtualQueryEx(m_proc_handle, addr, &mbi, sizeof(mbi))) return false;

			if(mbi.Protect & (PAGE_READWRITE/* | PAGE_WRITECOPY | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY | PAGE_WRITECOMBINE*/))
			{
				MemoryPage page;
				page.addr = addr;
				page.size = mbi.RegionSize;
				m_proc_pages.AddCpy(page);
			}
		}

		wxMessageBox(wxString::Format("%d 0x%p", m_proc_pages.GetCount(), GetBasicAddr()));

		return true;
	}

	u32 GetProcessCount()
	{
		return m_process_list.GetCount();
	}

	const PROCESSENTRY32& GetProcess(u32 n) const
	{
		return m_process_list[n];
	}

	const PROCESSENTRY32* GetProcess(const wxString& process) const
	{
		for(u32 i=0; i<m_process_list.GetCount(); ++i)
		{
			if(process.Cmp(m_process_list[i].szExeFile) != 0) continue;

			return &m_process_list[i];
		}

		return nullptr;
	}

	bool Search(Array<PAddr>& ret, const void* data, u32 types)
	{
		m_buf_size = block_size;
		m_buf = new u8[m_buf_size];

		for(u32 i=0; i<GetPageCount(); ++i)
		{
			SearchInPage(m_proc_pages[i], ret, data, types);
		}

		m_buf_size = 0;
		free(m_buf);

		return true;
	}

	bool Write(void* addr, const void* src, u32 size)
	{
		ULONG nwritten;

		if(!WriteProcessMemory(m_proc_handle, addr, src, size, &nwritten)) return false;

		return size == nwritten;
	}

	bool Read(void* addr, void* dst, u32 size)
	{
		ULONG nreaded;

		if(!ReadProcessMemory(m_proc_handle, addr, dst, size, &nreaded)) return false;

		return size == nreaded;
	}

	bool Write(const PAddr& addr)
	{
		return Write(addr.addr, &addr.value, addr.size);
	}

	bool Read(PAddr& addr)
	{
		addr.value = 0;

		return Read(addr.addr, &addr.value, addr.size);
	}

	void Close()
	{
		CloseHandle(m_proc_handle);
		m_proc_pages.Clear();
	}

	u8* GetBasicAddr()
	{
		return m_proc_pages.GetCount() ? m_proc_pages[0].addr : nullptr;
	}

private:
	u8* m_buf;
	u32 m_buf_size;

	u32 GetPageCount() const
	{
		return m_proc_pages.GetCount();
	}

	static const u32 block_size = 128;

	bool SearchInPage(const MemoryPage& page, Array<PAddr>& ret, const void* data, u32 types)
	{
		ULONG nreaded = 0;
		for(u8* addr = page.addr; addr<page.addr + page.size; addr += nreaded)
		{
			if(!ReadProcessMemory(m_proc_handle, addr, m_buf, m_buf_size, &nreaded))
			{
				//wxMessageBox("Search in page error");
			}

			if(!nreaded)
			{
				++addr;
				continue;
			}
			
			for(u32 t=0; t<m_types_count; ++t)
			{
				if(!(types & (1 << t)) || nreaded < m_types[t].size) continue;

				for(u32 i=0; i<=nreaded - m_types[t].size; ++i)
				{
					u64 cur_data = re(*(const u64*)data, m_types[t].size);
					if(memcmp(&cur_data, m_buf + i, m_types[t].size) != 0) continue;
					
					ret.Add(new PAddr(m_types[t], 1 << t, addr + i, m_types[t].size));
				}
			}
		}

		return true;
	}
};